Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[WIP] Stable Diffusion 3.x and Flux Optimization #22986

Draft
wants to merge 21 commits into
base: main
Choose a base branch
from
Draft

[WIP] Stable Diffusion 3.x and Flux Optimization #22986

wants to merge 21 commits into from

Conversation

tianleiwu
Copy link
Contributor

@tianleiwu tianleiwu commented Dec 2, 2024 •
edited
Loading

Description

This work is in progress.

Optimize the ONNX pipeline for Stable Diffusion 3.x and Flux 1.0 models (fp32 or fp16).

  • Update optimize_pipeline script
  • Update benchmkark script
  • Update document about Stable Diffusion 3.x and Flux 1.0 models
  • Add graph optimizations for MMDit model
    • FastGelu fusion
    • RMSNorm fusion
    • MultiHeadAttention fusion
  • Add graph optimizations for Flux transformer models
    • MultiHeadAttention fusion
  • Float16 conversion supports bfloat16 for blocked nodes.
  • Update graph optimizations for t5

Optimize the ONNX pipeline for Stable Diffusion 3.x and Flux 1.0 models:

python optimize_pipeline.py -i ./flux1_schnell_onnx/fp32 -o ./flux1_schnell_onnx/fp16 --float16

  Optimize flux1_schnell_onnx/fp32/transformer/model.onnx ...
  Fused LayerNormalization: 115
  Fused SimplifiedLayerNormalization: 152
  Fused FastGelu: 76
  Fused MultiHeadAttention: 57

H100 Benchmark Results

  • GPU: NVIDIA H100 80GB HBM3
  • Image Size: 1024x1024
  • Batch Size: 1
Model Steps Precision Engine Latency (Seconds) GPU Memory (MB)
Flux 1.0 Dev 50 BF16 Torch 2.5.1 (compile) 8.198 37,603
Flux 1.0 Dev 50 FP16+BF16 Optimum (ORT) 10.762 41,469
Flux 1.0 Dev 50 FP16+FP32 Optimum (ORT) 10.891 43,545
Flux 1.0 Dev 50 BF16 Torch 2.5.1 (eager) 12.339 36,651
Flux 1.0 Schnell 4 BF16 Torch 2.5.1 (compile) 0.775 37,857
Flux 1.0 Schnell 4 FP16+BF16 Optimum (ORT) 0.931 41,433
Flux 1.0 Schnell 4 FP16+FP32 Optimum (ORT) 0.939 43,809
Flux 1.0 Schnell 4 BF16 Torch 2.5.1 (eager) 1.120 36,629
SD 3.5 Large 50 BF16 Torch 2.5.1 (compile) 7.466 32,217
SD 3.5 Large 50 FP16+BF16 Optimum (ORT) 10.275 36,609
SD 3.5 Large 50 FP16+FP32 Optimum (ORT) 10.283 36,729
SD 3.5 Large 50 BF16 Torch 2.5.1 (eager) 11.615 31,517
SD 3.5 Medium 50 BF16 Torch 2.5.1 (compile) 3.240 21,143
SD 3.5 Medium 50 FP16+BF16 Optimum (ORT) 4.799 25,097
SD 3.5 Medium 50 FP16+FP32 Optimum (ORT) 4.838 25,109
SD 3.5 Medium 50 BF16 Torch 2.5.1 (eager) 5.582 20,489

A100 Benchmark Results

  • GPU: A100-SXM4-80GB
  • Image Size: 1024x1024
  • Batch Size: 1
Model Steps Precision Engine Latency (Seconds) GPU Memory (MB)
Flux 1.0 Dev 50 BF16 Torch 2.5.1 (compile) 17.593 37,723
Flux 1.0 Dev 50 FP16+BF16 Optimum (ORT) 21.918 41,348
Flux 1.0 Dev 50 FP16+FP32 Optimum (ORT) 22.060 44,860
Flux 1.0 Dev 50 BF16 Torch 2.5.1 (eager) 24.267 36,847
Flux 1.0 Schnell 4 BF16 Torch 2.5.1 (compile) 1.627 37,881
Flux 1.0 Schnell 4 FP16+BF16 Optimum (ORT) 1.884 41,537
Flux 1.0 Schnell 4 FP16+FP32 Optimum (ORT) 1.902 44,858
Flux 1.0 Schnell 4 BF16 Torch 2.5.1 (eager) 2.162 36,831
SD 3.5 Large 50 BF16 Torch 2.5.1 (compile) 15.881 32,307
SD 3.5 Large 50 FP16+FP32 Optimum (ORT) 19.837 36,451
SD 3.5 Large 50 FP16+BF16 Optimum (ORT) 19.964 36,461
SD 3.5 Large 50 BF16 Torch 2.5.1 (eager) 22.477 31,513
SD 3.5 Medium 50 BF16 Torch 2.5.1 (compile) 6.476 21,341
SD 3.5 Medium 50 FP16+FP32 Optimum (ORT) 8.775 25,183
SD 3.5 Medium 50 BF16 Torch 2.5.1 (eager) 10.057 20,433

Future Works

  • Triton kernel for matrix multiplication and auto tuning.
  • FP8/Int8 quantization

Motivation and Context

SD 3.5 Architecture:
https://huggingface.co/stabilityai/stable-diffusion-3.5-medium/resolve/main/mmdit-x.png

@tianleiwu tianleiwu marked this pull request as draft December 3, 2024 19:19
github-advanced-security[bot]
github-advanced-security bot found potential problems Dec 5, 2024
View reviewed changes
onnxruntime/python/tools/transformers/fusion_fastgelu.py
@@ -358,3 +361,122 @@
self.nodes_to_add.append(fused_node)
self.node_name_to_graph_name[fused_node.name] = self.this_graph_name
return True

def fuse_4(self, tanh_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]:

Check notice

Code scanning / CodeQL

Explicit returns mixed with implicit (fall through) returns Note

Mixing implicit and explicit returns may indicate an error as implicit returns always return None.
github-actions[bot]
github-actions bot reviewed Dec 19, 2024
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
github-advanced-security[bot]
github-advanced-security bot found potential problems Dec 20, 2024
View reviewed changes
onnxruntime/python/tools/transformers/onnx_model_mmdit.py
Comment on lines +91 to +96
# if (options is None) or options.enable_skip_layer_norm:
# self.fuse_skip_simplified_layer_norm()
# self.fuse_skip_layer_norm()
# if (options is None) or options.enable_bias_skip_layer_norm:
# # Fuse SkipLayerNormalization and Add Bias before it.
# self.fuse_add_bias_skip_layer_norm()

Check notice

Code scanning / CodeQL

Commented-out code Note

This comment appears to contain commented-out code.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments

@ACinfr ACinfr ACinfr left review comments

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@tianleiwu @ACinfr